#pragma once

#include "Texture/RHITexture.h"

class PipelineTransfer;

class TextureTransfer
{
	friend class PipelineTransfer;

	RefCountPtr<RHITexture> mTexture;

	VkImageMemoryBarrier mImageMemoryBarrier;

	VkPipelineStageFlags mSrcStageMask;
	VkPipelineStageFlags mDstStageMask;

	VkImageMemoryBarrier GetImageMemoryBarrier();

public:

	TextureTransfer();

	TextureTransfer& SetTexture(RefCountPtr<RHITexture> texture) noexcept;

	TextureTransfer& SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept;

	TextureTransfer& SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept;

	TextureTransfer& SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept;

	TextureTransfer& SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept;

	TextureTransfer& SetOldLayout(VkImageLayout oldLayout) noexcept;

	TextureTransfer& SetNewLayout(VkImageLayout newLayout) noexcept;

	TextureTransfer& SetSrcQueueFamily(VkQueueFlagBits srcQueueFlag) noexcept;

	TextureTransfer& SetDstQueueFamily(VkQueueFlagBits dstQueueFlag) noexcept;

	TextureTransfer& SetSubresourceRange(VkImageSubresourceRange subresourceRange) noexcept;

	void Transfer(VkQueueFlagBits queueFlag = VK_QUEUE_GRAPHICS_BIT, bool isMultiple = false);

	void Transfer(VkCommandBuffer commandBuffer, bool isMultiple = false);
};



class BufferTransfer
{
	friend class PipelineTransfer;

	RefCountPtr<RHIBuffer> mBuffer;

	VkBufferMemoryBarrier mBufferMemoryBarrier;

	VkPipelineStageFlags mSrcStageMask;
	VkPipelineStageFlags mDstStageMask;

	VkBufferMemoryBarrier GetBufferMemoryBarrier();

public:

	BufferTransfer();

	BufferTransfer& SetBuffer(RefCountPtr<RHIBuffer> buffer) noexcept;

	BufferTransfer& SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept;

	BufferTransfer& SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept;

	BufferTransfer& SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept;

	BufferTransfer& SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept;

	BufferTransfer& SetSrcQueueFamily(VkQueueFlagBits srcQueueFlag) noexcept;

	BufferTransfer& SetDstQueueFamily(VkQueueFlagBits dstQueueFlag) noexcept;

	BufferTransfer& SetOffset(VkDeviceSize offset) noexcept;

	BufferTransfer& SetSize(VkDeviceSize size) noexcept;

	void Transfer(VkQueueFlagBits queueFlag = VK_QUEUE_GRAPHICS_BIT, bool isMultiple = false);

	void Transfer(VkCommandBuffer commandBuffer, bool isMultiple = false);
};


class MemoryTransfer
{
	friend class PipelineTransfer;

	VkMemoryBarrier mMemoryBarrier;

	VkPipelineStageFlags mSrcStageMask;
	VkPipelineStageFlags mDstStageMask;

	VkMemoryBarrier GetMemoryBarrier();

public:

	MemoryTransfer();

	MemoryTransfer& SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept;

	MemoryTransfer& SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept;

	MemoryTransfer& SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept;

	MemoryTransfer& SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept;

	void Transfer(VkQueueFlagBits queueFlag = VK_QUEUE_GRAPHICS_BIT);

	void Transfer(VkCommandBuffer commandBuffer);
};

class PipelineTransfer
{
	std::vector<VkMemoryBarrier> mMemoryBarriers;
	std::vector<VkBufferMemoryBarrier> mBufferMemoryBarriers;
	std::vector<VkImageMemoryBarrier> mImageMemoryBarriers;

	VkPipelineStageFlags mSrcStageMask;
	VkPipelineStageFlags mDstStageMask;

public:

	PipelineTransfer& SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept;

	PipelineTransfer& SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept;

	PipelineTransfer& AddMemory(MemoryTransfer& memoryTransfer);

	PipelineTransfer& AddBuffer(BufferTransfer& bufferTransfer);

	PipelineTransfer& AddTexture(TextureTransfer& textureTransfer);

	void Transfer(VkQueueFlagBits queueFlag = VK_QUEUE_GRAPHICS_BIT);

	void Transfer(VkCommandBuffer commandBuffer);
};
